{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Building Local LLM Agents for Blender with Ryzen&trade; AI\n",
    "## Interactive Workshop Notebook\n",
    "\n",
    "Welcome to this hands-on workshop notebook! In this interactive session, you'll learn how to build powerful AI agents that can control Blender using natural language, all running locally on your Ryzen AI hardware.\n",
    "\n",
    "## <img src=\"./img/blender-agent.png\">\n",
    "\n",
    "### What You'll Learn\n",
    "- How to use GAIA & Lemonade, AMD's comprehensive open-source toolkit for local AI development\n",
    "- How to leverage Ryzen AI's NPU/iGPU hybrid mode for maximum performance with minimal power consumption\n",
    "- How to setup a bidirectional communication bridge between LLMs and Blender using the Model Context Protocol\n",
    "- How to build an intelligent agent with multi-step planning capabilities for complex 3D modeling tasks\n",
    "\n",
    "### Key Features\n",
    "- **Local Processing**: All AI processing happens on your device - no cloud dependencies\n",
    "- **Hybrid Performance**: Leverages Ryzen AI's NPU and iGPU for efficient execution\n",
    "- **Natural Language Interface**: Control Blender app using plain English commands\n",
    "- **Interactive Learning**: Step-by-step instruction with code examples for immediate feedback\n",
    "- **Rich Observability**: Visualize agent reasoning, execution paths, and performance metrics in real-time with detailed telemetry\n",
    "- **Real-World Application**: Create a working agent that can build basic 3D models\n",
    "\n",
    "### Workshop Notebook Structure\n",
    "This notebook is divided into five main parts:\n",
    "1. Introduction to GAIA and Lemonade\n",
    "2. LLM Client Implementation\n",
    "3. MCP Server Setup\n",
    "4. Blender Agent Implementation\n",
    "5. Testing and Evaluation\n",
    "\n",
    "### Prerequisites\n",
    "- AMD Ryzen AI 300 series processor or above\n",
    "- Blender 4.3 or newer [link](https://www.blender.org/download/releases/4-4)\n",
    "- Lemonade Server [link](https://lemonade-server.ai)\n",
    "- GAIA agent framework [link](https://github.com/amd/gaia)\n",
    "\n",
    "This notebook is best recommended for an audience that has basic python knowledge and familiarity with 3D concepts, but this is not required.  Now let's get started with setting up your environment and building your first local LLM agent for Blender!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Once [GAIA repo](https://github.com/amd/gaia) has been cloned and opened in VSCode, perform the following installation steps:\n",
    "\n",
    "1. Open a command prompt in terminal (NOT powershell):\n",
    "    1. `View -> Terminal`\n",
    "1. Checkout latest Blender branch if you haven't already:\n",
    "    1. `git checkout blender`\n",
    "    1. `git pull`\n",
    "1. Create and activate a conda environment if needed:\n",
    "    1. `conda create -n gaiaenv python=3.10 -y`\n",
    "    1. `conda activate gaiaenv`\n",
    "1. Install dependencies, using the following command while inside the cloned GAIA repo:\n",
    "    1. `pip install -e .[blender]`\n",
    "1. Set the default python interpreter:\n",
    "    1. `View -> Command Palette... -> Python: Select Interpreter`\n",
    "    1. Select `Python 3.10.17 ('gaiaenv')`\n",
    "1. Setup the Jupyter notebook (`blender.ipynb`) environment:\n",
    "    1. In this notebook, click on `Select Kernel -> Python Environments`\n",
    "    1. Choose the environment: `gaiaenv (Python 3.10.17)`\n",
    "1. Start the Lemonade Server:\n",
    "    1. Minimize VSCode and go to the Desktop\n",
    "    1. Double-click on `lemonade-server` desktop icon\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <img src=\"../src/gaia/interface/img/gaia.ico\" alt=\"favicon\" width=\"48\" height=\"48\"> Verify GAIA Installation\n",
    "\n",
    "Verify GAIA installation by checking the version.\n",
    "\n",
    "**Note**: The first time you run a cell in the notebook, a popup will ask you to select a Python Interpreter, make sure to select the one you had created earlier (typically called `gaiaenv`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gaia import version\n",
    "print(f\"GAIA version: {version.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <img src=\"../data/img/favicon.ico\" alt=\"favicon\" width=\"48\" height=\"48\"> Start Lemonade Server\n",
    "\n",
    "- The Lemonade Server tool comes pre-installed, simply **start** the tool by double clicking on the desktop icon called *lemonade-server*. You should see the following output:\n",
    "    ```bash\n",
    "    Starting Lemonade Server...\n",
    "    INFO:     Started server process [19180]\n",
    "    INFO:     Waiting for application startup.\n",
    "    INFO:     Application startup complete.\n",
    "    INFO:     Uvicorn running on http://localhost:8000 (Press CTRL+C to quit)\n",
    "    ```\n",
    "- *(Optional)* You can also run the server manually by opening a command terminal and running the following command:\n",
    "    ```bash\n",
    "    lemonade-server serve --log-level trace\n",
    "    ```\n",
    "    **NOTE**: `--log-level trace` enables more observability, printing the incoming prompts and performance results from the server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "- Now lets run a simple query with Lemonade via an OpenAI API using the Llama 3.2 3B parameter model on the Ryzen AI Hybrid device (iGPU and NPU)\n",
    "\n",
    "    **NOTE:** For more info, you can find Lemonade server details [here](https://github.com/lemonade-sdk/lemonade?tab=readme-ov-file)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "llm = OpenAI(base_url=\"http://localhost:8000/api/v0\", api_key=\"none\")\n",
    "\n",
    "response = llm.completions.create(\n",
    "    model=\"Llama-3.2-3B-Instruct-Hybrid\",\n",
    "    prompt=\"What is the capital of the moon?\"\n",
    ")\n",
    "print(response.choices[0].text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets enable streaming of the response for a better user experience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm.completions.create(\n",
    "    model=\"Llama-3.2-3B-Instruct-Hybrid\",\n",
    "    prompt=\"What is the capital of the moon?\",\n",
    "    stream=True\n",
    ")\n",
    "\n",
    "# Print each chunk as it arrives\n",
    "for chunk in response:\n",
    "    print(chunk.choices[0].text, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Context Protocol (MCP) for Blender Integration\n",
    "\n",
    "The Model-Context Protocol provides a bidirectional bridge between LLMs and Blender through a lightweight client-server architecture:\n",
    "\n",
    "### Client Component (Python Library)\n",
    "- **Socket Communication**: Establishes TCP connections to the Blender server (default port 9876)\n",
    "- **Command Serialization**: Converts Python function calls into properly formatted JSON commands\n",
    "- **Error Handling**: Manages network exceptions and command execution failures\n",
    "- **Pythonic Interface**: Wraps low-level socket communication in clean, simple API calls\n",
    "\n",
    "### Server Component (Blender Add-on)\n",
    "- **Threaded Socket Server**: Runs inside Blender as a non-blocking background thread\n",
    "- **Command Execution Pipeline**: Safely transfers JSON commands from socket to Blender's main thread\n",
    "- **Context Management**: Ensures operations execute in proper 3D viewport context\n",
    "- **Primitive Command Set**: Implements core operations like `create_object`, `modify_object`, and `execute_code`\n",
    "- **Response Formatting**: Standardizes all operation results into consistent JSON structures\n",
    "\n",
    "**NOTE**: The Blender MCP Server-Client implementation is a modified version from the following GitHub project [here](https://github.com/ahujasid/blender-mcp). We gratefully acknowledge the work of the BlenderMCP project developers used in this workshop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Blender MCP Server Setup\n",
    "\n",
    "Lets install the Blender MCP server:\n",
    "1. Open Blender\n",
    "2. Go to `Edit > Preferences > Add-ons`\n",
    "3. Using the down arrow button in the top-right corner, click `Install...` and navigate to `<root>/src/gaia/mcp/blender_mcp_server.py`\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"./img/blender-plugin.png\" alt=\"favicon\">\n",
    "</p>\n",
    "\n",
    "4. Ensure the add-on is enabled by confirming that the check box next to `Simple Blender MCP` is selected\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"./img/blender-plugin2.png\" alt=\"favicon\">\n",
    "</p>\n",
    "\n",
    "  - **NOTE**: If plugin does not appear, try `Refresh Local`. If that does not work, then delete the plug-in from `C:\\Users\\<username>\\AppData\\Roaming\\Blender Foundation\\Blender\\4.4\\scripts\\addons` and add it again.\n",
    "\n",
    "5. Once installed, you'll find a new panel in the 3D viewport sidebar called `Blender MCP`. This panel allows you to:\n",
    "  - Start/Stop the MCP server\n",
    "  - Configure the port number (default: 9876)\n",
    "  - See the current server status\n",
    "\n",
    "    **NOTE**: If you do not see the viewport sidebar, hit `N` key or go to `View->Sidebar` and enable the checkbox.\n",
    "\n",
    "6. Click `Start Server` and make sure you are running on port `9876`\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"./img/blender-plugin3.png\" alt=\"favicon\">\n",
    "</p>\n",
    "\n",
    "7. Now lets test the MCP client-server connection using a simple socket client that can send and receive JSON messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from gaia.mcp.blender_mcp_client import MCPClient\n",
    "\n",
    "mcp = MCPClient(host='localhost', port=9876)\n",
    "pprint(mcp.get_scene_info())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Let's Make a Blue Sphere!\n",
    "\n",
    "This little test puts a blue sphere right above the default cube in Blender. It shows how easy it is to:\n",
    "\n",
    "- Connect to Blender with our MCP client\n",
    "- Create 3D objects with just a few lines of code\n",
    "- Add a nice blue material to make it pop\n",
    "\n",
    "Perfect for checking if your Blender connection works.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "# Create a sphere 2.5 units above the origin (just above the default cube)\n",
    "sphere = mcp.create_object(\n",
    "    type=\"SPHERE\", \n",
    "    name=\"Workshop_Sphere\", \n",
    "    location=(0, 0, 2.5),\n",
    "    scale=(0.75, 0.75, 0.75)  # Slightly smaller than default\n",
    ")\n",
    "\n",
    "pprint(sphere)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "material_code = \"\"\"\n",
    "import bpy\n",
    "# Create a blue material for the sphere\n",
    "mat = bpy.data.materials.new(name=\"Sphere_Material\")\n",
    "mat.diffuse_color = (0.1, 0.3, 0.8, 1.0)  # Blue color\n",
    "\n",
    "# Find our sphere and assign material\n",
    "sphere = bpy.data.objects.get(\"Workshop_Sphere\")\n",
    "if sphere:\n",
    "    if sphere.data.materials:\n",
    "        sphere.data.materials[0] = mat\n",
    "    else:\n",
    "        sphere.data.materials.append(mat)\n",
    "    result = \"Material applied to sphere!\"\n",
    "else:\n",
    "    result = \"Sphere not found!\"\n",
    "\"\"\"\n",
    "material_result = mcp.execute_code(material_code)\n",
    "\n",
    "print(material_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a Simple Blender Agent\n",
    "\n",
    "Let's create a basic agent that can understand natural language commands and create 3D objects in Blender!\n",
    "\n",
    "Step-by-Step Plan:\n",
    "1. Set up the LLM client to process natural language\n",
    "1. Set up the MCP client to communicate with Blender\n",
    "1. Create the agent to connect these components\n",
    "1. Test with simple commands\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Import required components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gaia.llm.llm_client import LLMClient\n",
    "from gaia.mcp.blender_mcp_client import MCPClient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Set up the LLM client with system prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are a 3D modeling assistant. IMPORTANT: For EACH user request, respond with EXACTLY ONE LINE in this format:\n",
    "TYPE,x,y,z,sx,sy,sz\n",
    "\n",
    "Where:\n",
    "- TYPE is one of: CUBE, SPHERE, CYLINDER, CONE, TORUS - no other types allowed\n",
    "- x,y,z are the LOCATION coordinates in 3D space (must be numbers)\n",
    "- sx,sy,sz are the SCALE factors (must be numbers)\n",
    "\n",
    "You MUST include ALL 7 parameters separated by commas.\n",
    "You MUST respond with ONLY ONE LINE.\n",
    "You MUST NOT include any other text.\n",
    "\n",
    "Example: \"Create a large sphere at the origin\" → SPHERE,0,0,0,2,2,2\n",
    "Example: \"Make a tall cylinder\" → CYLINDER,0,0,0,1,1,3\n",
    "\"\"\"\n",
    "\n",
    "llm_client = LLMClient(\n",
    "    use_local=True,\n",
    "    system_prompt=system_prompt,\n",
    "    base_url=\"http://localhost:8000/api/v0\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Set up the MCP client for Blender"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "mcp_client = MCPClient(host='localhost', port=9876)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Create a simple function to process commands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_object_from_description(description):\n",
    "    # Get object parameters from LLM\n",
    "    query = description.strip()\n",
    "    llm_response = llm_client.generate(query)\n",
    "    print(f\"\\nQuery: {query}\")\n",
    "    print(f\"LLM response: {llm_response}\")\n",
    "\n",
    "    # Parse the response\n",
    "    try:\n",
    "        parts = llm_response.split(',')\n",
    "        if len(parts) != 7:\n",
    "            raise ValueError(f\"Expected 7 parts in response, got {len(parts)}\")\n",
    "            \n",
    "        obj_type = parts[0].strip().upper()\n",
    "        location = (float(parts[1]), float(parts[2]), float(parts[3]))\n",
    "        scale = (float(parts[4]), float(parts[5]), float(parts[6]))\n",
    "\n",
    "        # Create the object in Blender\n",
    "        result = mcp_client.create_object(\n",
    "            type=obj_type,\n",
    "            name=f\"nlp_{obj_type.lower()}\",\n",
    "            location=location,\n",
    "            scale=scale\n",
    "        )\n",
    "\n",
    "        print(f\"✓ Created {obj_type} at location {location} with scale {scale}\")\n",
    "        return result\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Error: {str(e)}\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test Your Agent\n",
    "\n",
    "Let's try some examples to see if our agent can understand natural language commands:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\n",
    "    \"Create a cube at the origin\",\n",
    "    \"Make a tall skinny cylinder that's 3 units above the origin\",\n",
    "    \"Put a large flat sphere 2 units to the right of the origin\",\n",
    "    \"Create a small cone in front of the cube\"\n",
    "]\n",
    "\n",
    "for example in examples:\n",
    "    create_object_from_description(example)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the basic agent is working, lets make it more robust by allowing it to iterate and correct its output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a Blender Agent: Step-by-Step Guide\n",
    "\n",
    "### Introduction: The Base Agent Framework\n",
    "The Gaia Agent system provides a powerful foundation for building specialized AI agents. At its core, the Agent framework handles:\n",
    "\n",
    "- **Tool Registration**: Converts Python functions into AI-callable tools with proper argument handling, documentation, and error management\n",
    "- **LLM Integration**: Manages communication with large language models through a unified interface\n",
    "- **Conversation Planning**: Implements sophisticated planning with multi-step execution\n",
    "- **Error Resilience**: Built-in error recovery mechanisms for adapting to failures\n",
    "- **Rich Observability**: Comprehensive visibility into agent reasoning and execution\n",
    "\n",
    "Our BlenderAgent extends this framework by adding 3D modeling capabilities through specialized tools that communicate with Blender's Python API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Import and Setup\n",
    "\n",
    "*NOTE: The decorator pattern is perfect for notebooks - it lets us build our agent incrementally, adding methods in different cells while keeping code organization clean and intuitive.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import base Agent components and Blender communication\n",
    "from gaia.agents.base.agent import Agent\n",
    "from gaia.agents.base.tools import tool\n",
    "from gaia.agents.base.console import AgentConsole\n",
    "from gaia.mcp.blender_mcp_client import MCPClient\n",
    "import logging\n",
    "\n",
    "# Set up logging\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Create our method decorator for building the class incrementally\n",
    "def add_method(cls):\n",
    "    \"\"\"\n",
    "    Decorator that adds a function as a method to a class.\n",
    "    Perfect for incrementally building out functionality in a notebook.\n",
    "    \"\"\"\n",
    "    def decorator(func):\n",
    "        setattr(cls, func.__name__, func)\n",
    "        return func\n",
    "    return decorator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Create the BlenderAgent Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create our BlenderAgent class extending the base Agent\n",
    "class BlenderAgent(Agent):\n",
    "    \"\"\"\n",
    "    Blender-specific agent focused on 3D scene creation and modification.\n",
    "    Inherits core functionality from the base Agent class.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        use_local_llm=True,\n",
    "        mcp=None,\n",
    "        model_id=None,\n",
    "        base_url=\"http://localhost:8000/api/v0\",\n",
    "        max_steps=5,\n",
    "        debug_prompts=False,\n",
    "        output_dir=None,\n",
    "        streaming=False,\n",
    "        show_stats=True\n",
    "    ):\n",
    "        \"\"\"Initialize the BlenderAgent with MCP client and LLM client\"\"\"\n",
    "        # Initialize the MCP client for Blender communication\n",
    "        self.mcp = mcp if mcp else MCPClient()\n",
    "        \n",
    "        # Call the parent class constructor for LLM integration\n",
    "        super().__init__(\n",
    "            use_local_llm=use_local_llm,\n",
    "            model_id=model_id,\n",
    "            base_url=base_url,\n",
    "            max_steps=max_steps,\n",
    "            debug_prompts=debug_prompts,\n",
    "            output_dir=output_dir,\n",
    "            streaming=streaming,\n",
    "            show_stats=show_stats\n",
    "        )\n",
    "\n",
    "    def _create_console(self):\n",
    "        return AgentConsole()\n",
    "\n",
    "    # Stub implementations for all abstract methods,\n",
    "    # will be updated further down.\n",
    "    def _get_system_prompt(self):\n",
    "        # Stub implementation \n",
    "        return \"\"\n",
    "    \n",
    "    def _register_tools(self):\n",
    "        # Stub implementation\n",
    "        pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Implement System Prompt\n",
    "\n",
    "These methods fulfill the abstract requirements of the base Agent class. The system prompt is particularly important - it shapes how the LLM approaches tasks, enforcing proper planning and workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "@add_method(BlenderAgent)\n",
    "def _get_system_prompt(self) -> str:\n",
    "    \"\"\"Generate the system prompt for the Blender agent.\"\"\"\n",
    "    return \"\"\"\n",
    "You are a specialized Blender 3D assistant that can create and modify 3D scenes.\n",
    "You will use a set of tools to accomplish tasks based on the user's request.\n",
    "\n",
    "==== JSON RESPONSE FORMAT ====\n",
    "ALWAYS respond with a single valid JSON object. NO text outside this structure.\n",
    "- Use double quotes for keys and string values\n",
    "- Ensure all braces and brackets are properly closed\n",
    "- No trailing commas in arrays or objects\n",
    "- All required fields must be included\n",
    "- Never wrap your JSON in code blocks or backticks\n",
    "\n",
    "Your JSON response must follow this format:\n",
    "{\n",
    "    \"thought\": \"your reasoning about what to do\",\n",
    "    \"goal\": \"clear statement of what you're achieving\",\n",
    "    \"plan\": [\n",
    "        {\"tool\": \"tool1\", \"tool_args\": {\"arg1\": \"val1\"}},\n",
    "        {\"tool\": \"tool2\", \"tool_args\": {\"arg1\": \"val1\"}}\n",
    "    ],\n",
    "    \"tool\": \"first_tool_to_execute\",\n",
    "    \"tool_args\": {\"arg1\": \"val1\", \"arg2\": \"val2\"}\n",
    "}\n",
    "\n",
    "For final answers:\n",
    "{\n",
    "    \"thought\": \"your reasoning\",\n",
    "    \"goal\": \"what was achieved\",\n",
    "    \"answer\": \"your final answer\"\n",
    "}\n",
    "\n",
    "==== CRITICAL RULES ====\n",
    "1. ALWAYS create a plan before executing any tools\n",
    "2. Each plan step must be atomic (one tool call per step)\n",
    "3. For colored objects, ALWAYS include both create_object AND set_material_color steps\n",
    "4. When clearing a scene, ONLY use clear_scene without creating new objects unless requested\n",
    "5. Always use the actual returned object names for subsequent operations\n",
    "6. Never repeat the same tool call with identical arguments\n",
    "\n",
    "==== COMMON WORKFLOWS ====\n",
    "1. Clearing a scene: Use clear_scene() with no arguments\n",
    "2. Creating a colored object: First create_object, then set_material_color\n",
    "3. Modifying objects: Use modify_object with the parameters you want to change\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Register Available Tools\n",
    "\n",
    "The *@tool* decorator automatically registers these functions with the base Agent's tool registry. These tools provide the basic building blocks for 3D scene creation - clearing the scene and creating primitive 3D objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, Any, Optional\n",
    "\n",
    "@add_method(BlenderAgent)\n",
    "def _register_tools(self):\n",
    "    \"\"\"Register all Blender-related tools for the agent.\"\"\"\n",
    "\n",
    "    @tool\n",
    "    def clear_scene() -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Remove all objects from the current Blender scene.\n",
    "\n",
    "        Returns:\n",
    "            Dictionary containing the operation result\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will clear the scene to start fresh\",\n",
    "            \"goal\": \"Clear the scene to start fresh\",\n",
    "            \"tool\": \"clear_scene\",\n",
    "            \"tool_args\": {}\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            from gaia.agents.Blender.core.scene import SceneManager\n",
    "            scene_manager = SceneManager(self.mcp)\n",
    "            return scene_manager.clear_scene()\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "\n",
    "    @tool\n",
    "    def create_object(type: str = \"CUBE\", name: str = None, location: tuple = (0,0,0), rotation: tuple = (0,0,0), scale: tuple = (1,1,1)) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Create a 3D object in Blender.\n",
    "\n",
    "        Args:\n",
    "            type: Object type (CUBE, SPHERE, CYLINDER, CONE, TORUS)\n",
    "            name: Optional name for the object (default: generated from type)\n",
    "            location: (x, y, z) coordinates for object position (default: (0,0,0))\n",
    "            rotation: (rx, ry, rz) rotation in radians (default: (0,0,0))\n",
    "            scale: (sx, sy, sz) scaling factors for the object (default: (1,1,1))\n",
    "\n",
    "        Returns:\n",
    "            Dictionary containing the creation result\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will create a cube at the center of the scene\",\n",
    "            \"goal\": \"Create a red cube at the center of the scene\",\n",
    "            \"tool\": \"create_object\",\n",
    "            \"tool_args\": {\n",
    "                \"type\": \"CUBE\",\n",
    "                \"name\": \"my_cube\",\n",
    "                \"location\": [0, 0, 0],\n",
    "                \"rotation\": [0, 0, 0],\n",
    "                \"scale\": [1, 1, 1]\n",
    "            }\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            print(f\"create_object: {type}, {name}, {location}, {rotation}, {scale}\")\n",
    "            result = self.mcp.create_object(\n",
    "                type=type.upper(),\n",
    "                name=name or f\"generated_{type.lower()}\",\n",
    "                location=location,\n",
    "                rotation=rotation,\n",
    "                scale=scale\n",
    "            )\n",
    "            return result\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    @tool\n",
    "    def set_material_color(object_name: str, color: tuple = (1, 0, 0, 1)) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Set the material color for an object. Creates a new material if one doesn't exist.\n",
    "\n",
    "        Args:\n",
    "            object_name: Name of the object to modify\n",
    "            color: RGBA color values as tuple (red, green, blue, alpha), values from 0-1\n",
    "\n",
    "        Returns:\n",
    "            Dictionary with the operation result\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will set the cube's material to red\",\n",
    "            \"goal\": \"Create a red cube at the center of the scene\",\n",
    "            \"tool\": \"set_material_color\",\n",
    "            \"tool_args\": {\n",
    "                \"object_name\": \"my_cube\",\n",
    "                \"color\": [1, 0, 0, 1]\n",
    "            }\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            from gaia.agents.Blender.core.materials import MaterialManager\n",
    "            material_manager = MaterialManager(self.mcp)\n",
    "            return material_manager.set_material_color(object_name, color)\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    # @tool\n",
    "    def get_object_info(name: str) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Get information about an object in the scene.\n",
    "\n",
    "        Args:\n",
    "            name: Name of the object\n",
    "\n",
    "        Returns:\n",
    "            Dictionary containing object information\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will get information about the cube\",\n",
    "            \"goal\": \"Create a red cube at the center of the scene\",\n",
    "            \"tool\": \"get_object_info\",\n",
    "            \"tool_args\": {\n",
    "                \"name\": \"my_cube\"\n",
    "            }\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            return self.mcp.get_object_info(name)\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    @tool\n",
    "    def modify_object(name: str, location: tuple = None, scale: tuple = None, rotation: tuple = None) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Modify an existing object in Blender.\n",
    "\n",
    "        Args:\n",
    "            name: Name of the object to modify\n",
    "            location: New (x, y, z) location or None to keep current\n",
    "            scale: New (sx, sy, sz) scale or None to keep current\n",
    "            rotation: New (rx, ry, rz) rotation or None to keep current\n",
    "\n",
    "        Returns:\n",
    "            Dictionary with the modification result\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will move the cube up by 2 units\",\n",
    "            \"goal\": \"Create a red cube at the center of the scene\",\n",
    "            \"tool\": \"modify_object\",\n",
    "            \"tool_args\": {\n",
    "                \"name\": \"my_cube\",\n",
    "                \"location\": [0, 0, 2],\n",
    "                \"scale\": null,\n",
    "                \"rotation\": null\n",
    "            }\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            return self.mcp.modify_object(\n",
    "                name=name,\n",
    "                location=location,\n",
    "                scale=scale,\n",
    "                rotation=rotation\n",
    "            )\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    # @tool\n",
    "    def delete_object(name: str) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Delete an object from the scene.\n",
    "\n",
    "        Args:\n",
    "            name: Name of the object to delete\n",
    "\n",
    "        Returns:\n",
    "            Dictionary with the deletion result\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will delete the cube\",\n",
    "            \"goal\": \"Clear the scene to start fresh\",\n",
    "            \"tool\": \"delete_object\",\n",
    "            \"tool_args\": {\n",
    "                \"name\": \"my_cube\"\n",
    "            }\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            return self.mcp.delete_object(name)\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    @tool\n",
    "    def get_scene_info() -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Get information about the current scene.\n",
    "\n",
    "        Returns:\n",
    "            Dictionary containing scene information\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will get information about the current scene\",\n",
    "            \"goal\": \"Clear the scene to start fresh\",\n",
    "            \"tool\": \"get_scene_info\",\n",
    "            \"tool_args\": {}\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            return self.mcp.get_scene_info()\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    # @tool\n",
    "    def execute_blender_code(code: str) -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Execute arbitrary Python code in Blender with error handling.\n",
    "\n",
    "        Args:\n",
    "            code: Python code to execute in Blender\n",
    "\n",
    "        Returns:\n",
    "            Dictionary with execution results or error information\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will execute custom code to create a complex shape\",\n",
    "            \"goal\": \"Create a red cube at the center of the scene\",\n",
    "            \"tool\": \"execute_blender_code\",\n",
    "            \"tool_args\": {\n",
    "                \"code\": \"import bpy\\\\nbpy.ops.mesh.primitive_cube_add()\"\n",
    "            }\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            return self.mcp.execute_code(code)\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}\n",
    "\n",
    "    # @tool\n",
    "    def diagnose_scene() -> Dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Diagnose the current Blender scene for common issues.\n",
    "        Returns information about objects, materials, and potential problems.\n",
    "\n",
    "        Returns:\n",
    "            Dictionary with diagnostic information\n",
    "\n",
    "        Example JSON response:\n",
    "        ```json\n",
    "        {\n",
    "            \"thought\": \"I will diagnose the scene for any issues\",\n",
    "            \"goal\": \"Clear the scene to start fresh\",\n",
    "            \"tool\": \"diagnose_scene\",\n",
    "            \"tool_args\": {}\n",
    "        }\n",
    "        ```\n",
    "        \"\"\"\n",
    "        try:\n",
    "            # Use the core library's scene diagnosis code generator\n",
    "            diagnostic_code = generate_scene_diagnosis_code()\n",
    "            return self.mcp.execute_code(diagnostic_code)\n",
    "        except Exception as e:\n",
    "            self.error_history.append(str(e))\n",
    "            return {\"status\": \"error\", \"error\": str(e)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Add Object Name Tracking\n",
    "\n",
    "This intelligent name-tracking mechanism is critical for multi-step operations. It detects when Blender alters object names and automatically updates future steps in the plan, enabling complex workflows to succeed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "@add_method(BlenderAgent)\n",
    "def _post_process_tool_result(self, tool_name: str, tool_args: Dict[str, Any], tool_result: Dict[str, Any]) -> None:\n",
    "    \"\"\"\n",
    "    Post-process the tool result for Blender-specific handling.\n",
    "\n",
    "    Args:\n",
    "        tool_name: Name of the tool that was executed\n",
    "        tool_args: Arguments that were passed to the tool\n",
    "        tool_result: Result returned by the tool\n",
    "    \"\"\"\n",
    "    # Track object name if created\n",
    "    if tool_name == \"create_object\":\n",
    "        actual_name = self._track_object_name(tool_result)\n",
    "        if actual_name:\n",
    "            logger.debug(f\"Actual object name created: {actual_name}\")\n",
    "            self.console.print_info(f\"Note: Blender assigned name '{actual_name}' to the created object\")\n",
    "\n",
    "            # Update subsequent steps in the plan that might use this object\n",
    "            if (self.current_plan and self.current_step < len(self.current_plan) - 1):\n",
    "                for i in range(self.current_step + 1, len(self.current_plan)):\n",
    "                    future_step = self.current_plan[i]\n",
    "                    if isinstance(future_step, dict) and \"tool_args\" in future_step:\n",
    "                        args = future_step[\"tool_args\"]\n",
    "                        # Look for object_name or name parameters\n",
    "                        if \"object_name\" in args and args[\"object_name\"] == tool_args.get(\"name\"):\n",
    "                            logger.debug(f\"Updating object_name in future step {i+1} from {args['object_name']} to {actual_name}\")\n",
    "                            self.current_plan[i][\"tool_args\"][\"object_name\"] = actual_name\n",
    "                        if \"name\" in args and args[\"name\"] == tool_args.get(\"name\"):\n",
    "                            logger.debug(f\"Updating name in future step {i+1} from {args['name']} to {actual_name}\")\n",
    "                            self.current_plan[i][\"tool_args\"][\"name\"] = actual_name\n",
    "\n",
    "def _track_object_name(self, result):\n",
    "    \"\"\"\n",
    "    Extract and track the actual object name returned by Blender.\n",
    "\n",
    "    Args:\n",
    "        result: The result dictionary from a tool execution\n",
    "\n",
    "    Returns:\n",
    "        The actual object name if found, None otherwise\n",
    "    \"\"\"\n",
    "    try:\n",
    "        if isinstance(result, dict):\n",
    "            if result.get('status') == 'success':\n",
    "                if 'result' in result and isinstance(result['result'], dict):\n",
    "                    # Extract name from create_object result\n",
    "                    if 'name' in result['result']:\n",
    "                        actual_name = result['result']['name']\n",
    "                        logger.debug(f\"Extracted object name: {actual_name}\")\n",
    "                        return actual_name\n",
    "        return None\n",
    "    except Exception as e:\n",
    "        logger.error(f\"Error extracting object name: {str(e)}\")\n",
    "        return None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6: Add Scene Creation Method\n",
    "\n",
    "This method showcases how we leverage the base Agent's planning and execution capabilities for complex tasks. It allocates more steps for scene creation and properly formats the query for optimal LLM understanding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@add_method(BlenderAgent)\n",
    "def create_interactive_scene(\n",
    "    self,\n",
    "    scene_description,\n",
    "    max_steps=None,\n",
    "    output_to_file=True,\n",
    "    filename=None\n",
    "):\n",
    "    \"\"\"\n",
    "    Create a complex scene from a natural language description.\n",
    "    \n",
    "    The base Agent's planning capabilities shine here, intelligently\n",
    "    breaking down complex scenes into sequences of atomic operations.\n",
    "    \n",
    "    Args:\n",
    "        scene_description: Description of the scene to create\n",
    "        max_steps: Maximum number of steps (defaults to twice the standard limit)\n",
    "        output_to_file: If True, write results to a JSON file\n",
    "        filename: Optional filename for output\n",
    "        \n",
    "    Returns:\n",
    "        Dict containing the scene creation result\n",
    "    \"\"\"\n",
    "    # Use the process_query method from the base Agent\n",
    "    return self.process_query(\n",
    "        f\"Create a complete 3D scene with the following description: {scene_description}\",\n",
    "        max_steps=max_steps if max_steps is not None else self.max_steps * 2,\n",
    "        output_to_file=output_to_file,\n",
    "        filename=filename\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7: Initialize and Print Tools\n",
    "\n",
    "This initialization calls our additional tool registration method and creates a fully-functional agent instance. The base Agent framework handles all the LLM communication, planning, and execution orchestration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update class with all enhanced methods\n",
    "BlenderAgent._get_system_prompt = _get_system_prompt\n",
    "BlenderAgent._register_tools = _register_tools\n",
    "BlenderAgent._post_process_tool_result = _post_process_tool_result\n",
    "BlenderAgent._track_object_name = _track_object_name\n",
    "\n",
    "# Create a BlenderAgent instance\n",
    "agent = BlenderAgent(\n",
    "    model_id=\"Llama-3.2-3B-Instruct-Hybrid\",\n",
    "    max_steps=5,\n",
    "    output_dir=\"./output\",\n",
    "    streaming=False\n",
    ")\n",
    "\n",
    "agent.list_tools()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 8: Test the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test with a basic example\n",
    "print(\"\\n🔍 TESTING: Creating a colored object\\n\")\n",
    "result = agent.process_query(\"Create a red cube at the center of the scene and make sure it has a red material\")\n",
    "\n",
    "# Display the agent's structured plan and execution results\n",
    "agent.display_result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 9: Create a Complex Scene\n",
    "\n",
    "For complex scenes, the agent leverages the LLM's planning to decompose the description into a sequence of coordinated operations. The base Agent handles plan execution, adjusting to actual object names and recovering from any errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a more complex scene with a complete description\n",
    "print(\"\\n🏙️ CREATING: A complete desk scene\\n\")\n",
    "\n",
    "# Create the scene with up to 10 steps allowed\n",
    "scene_result = agent.create_interactive_scene(\n",
    "    \"Create a wooden desk with a computer monitor, keyboard, and a blue coffee mug on it\",\n",
    "    max_steps=10\n",
    ")\n",
    "\n",
    "# Display the complete execution results\n",
    "agent.display_result()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
